.github/workflows/integration_tests.yml (1)

61-64: CI env echo is fine; simpler shell is enough.

You can avoid spinning up Python here.

Apply:

-      - name: Is running on CI environment (GitHub Actions)?
-        run: |
-          python -c "import os; print('Result: ', os.getenv('GITHUB_ACTIONS', 'Not set'))"
+      - name: Is running on CI environment (GitHub Actions)?
+        run: echo "Result: ${GITHUB_ACTIONS:-Not set}"

src/midst_toolkit/models/clavaddpm/synthesizer.py (1)

790-793: Directory creation before persisting: good; prefer explicit pickle protocol.

Use highest protocol for speed/size.

-        with open(before_matching_dir / "synthetic_tables.pkl", "wb") as file:
-            pickle.dump(synthetic_tables, file)
+        with open(before_matching_dir / "synthetic_tables.pkl", "wb") as file:
+            pickle.dump(synthetic_tables, file, protocol=pickle.HIGHEST_PROTOCOL)

tests/integration/utils.py (1)

4-11: CI detection works; small robustness tweak.

Make it case-insensitive to avoid surprises.

-    return os.getenv("GITHUB_ACTIONS", "false") == "true"
+    return os.getenv("GITHUB_ACTIONS", "false").lower() == "true"

tests/integration/models/clavaddpm/test_synthesizer.py (2)

66-66: Minor: fix test name typo.

Keeps repo tidy and aids discovery.

-def test_clava_syntheesize_multi_table(tmp_path: Path):
+def test_clava_synthesize_multi_table(tmp_path: Path):

---

73-86: Unpack training return to improve readability and ensure updated tables are used.

The return type is tuple[Tables, dict[...]]. Unpacking avoids tuple indexing and ensures the updated tables from training (with added placeholder column for root tables) is passed to synthesis, rather than the original tables.

-    models = clava_training(tables, relation_order, tmp_path, DIFFUSION_CONFIG, CLASSIFIER_CONFIG, device=DEVICE)
+    tables, models = clava_training(tables, relation_order, tmp_path, DIFFUSION_CONFIG, CLASSIFIER_CONFIG, device=DEVICE)
@@
-        models[1],
+        models,

The DEVICE variable already correctly falls back to "cpu" on systems without CUDA support.

tests/integration/models/clavaddpm/test_model.py (3)

293-307: Consider using pytest warnings for better visibility.

The CI-aware testing approach is reasonable for handling platform differences. However, logging.warning may not be visible in test output depending on the test runner configuration.

Consider using pytest's warning mechanism for better visibility:

-        logging.warning("Not running on CI, assertions are made with a higher tolerance.")
+        pytest.warns(UserWarning, match="Not running on CI, assertions are made with a higher tolerance.")

Or use pytest.skip with a reason if the assertions should truly be skipped on non-CI environments:

    if not is_running_on_ci_environment():
        pytest.skip("Strict assertions only run on CI due to platform differences")

---

350-365: CI-aware assertions follow consistent pattern.

Moving model_layers extraction before the conditional block avoids duplication. The same recommendation about using pytest warnings (from lines 293-307) applies here as well.

---

385-389: Conditional sample assertions completely skipped on non-CI.

Unlike previous CI-aware checks that use looser tolerances, this block completely skips assertions on non-CI environments, reducing test coverage.

Consider adding at least a basic sanity check on non-CI:

     else:
         logging.warning("Not running on CI, skipping detailed assertions.")
+        # Basic sanity checks even on non-CI
+        assert conditional_sample.shape == expected_conditional_sample.shape
+        assert not torch.isnan(conditional_sample).any()
+        assert not torch.isinf(conditional_sample).any()

This ensures the operation doesn't produce invalid results even if exact values differ across platforms.

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

tests/integration/models/clavaddpm/test_synthesizer.py (1)

91-96: Acknowledged: pickle usage flagged in past review.

This was already noted in previous review comments. The pickle.loads pattern is acceptable for repository-controlled test fixtures but carries security risks with untrusted data.

Based on past review comments, consider refactoring to use pickle.load with a file handle:

-        expected_cleaned_tables = pickle.loads(
-            Path("tests/integration/assets/multi_table/assertion_data/cleaned_tables.pkl").read_bytes(),
-        )
+        with open("tests/integration/assets/multi_table/assertion_data/cleaned_tables.pkl", "rb") as f:
+            expected_cleaned_tables = pickle.load(f)

tests/integration/models/clavaddpm/test_synthesizer.py (1)

72-72: Clarify the tuple unpacking for better readability.

The code works correctly but is confusing: models at line 72 receives the full tuple (tables, models_dict) from clava_training, then line 83 uses models[1] to extract the models dictionary. This makes the variable name models misleading.

Apply this diff for clearer code:

-    models = clava_training(tables, relation_order, tmp_path, DIFFUSION_CONFIG, CLASSIFIER_CONFIG, device=DEVICE)
+    tables, models = clava_training(tables, relation_order, tmp_path, DIFFUSION_CONFIG, CLASSIFIER_CONFIG, device=DEVICE)
 
     # TODO: Temporary, we should refactor those configs
     configs = deepcopy(SYNTHESIZING_CONFIG)
     configs["general"]["workspace_dir"] = str(tmp_path)
 
     cleaned_tables, _, _ = clava_synthesizing(
         tables,
         relation_order,
         tmp_path,
         all_group_lengths_prob_dicts,
-        models[1],
+        models,
         configs,
     )

Also applies to: 78-85

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro